package markov;


import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

import everything.WorldMap;
import bnetwork.ClimateNetwork;
import bnetwork.Node;

public class MDP {
	//for each flag - flag values (0 or 1)
	//for each point - ice or no ice

	//actions
	
	static final float gamma = 0.9f;
	
	public MDP(){
		
	}
	
	public MDPStateTable[] compute(ClimateNetwork n, WorldMap m, int numMoves, boolean remember){
		int numFlags = m.getNumFlags();
		int numIces = m.getNumIces();
		int flagsPow =	(int)Math.pow(2, numFlags);
		int icesPow = (int)Math.pow(3, numIces);
		int height = m.getHeight();
		int width = m.getWidth();
		int stateSpace = flagsPow*icesPow*height*width*5-1;
		float[][] iceProbs = new float[icesPow][numIces];
		
		System.out.println("Computing all possible probabilities");
		for (int i=0; i<iceProbs.length; i++){
			Map<Node, Integer> newEvidence = getIceEvidence(i, numIces, m, n);
			newEvidence.putAll(n.getEvidence());
			
			for (int j=0; j<numIces; j++){
				int nodeId = m.getIceCellId(j);
				iceProbs[i][j] = n.getProbIce(newEvidence, nodeId);
			}
		}
		System.out.println("Done computing");
		
		MDPStateTable[] tables = new MDPStateTable[numMoves];
		//fill up first table - with states and values of 0 - the value is just the reward for each state.
		MDPStateTable startTable;
		if (remember){
			tables[0] = new MDPStateTable(stateSpace);
			startTable = tables[0];
		} else {
			tables[1] = new MDPStateTable(stateSpace);
			startTable = tables[1];
		}
		System.out.println("State space is " + stateSpace + " big");
		System.out.println("Step: " + 0);
		for (int i=0; i<stateSpace; i++){
			MDPState p = new MDPState(i, numFlags, numIces, height, width);
			if (m.brickStatus(p.xPosition, p.yPosition)==WorldMap.WALL){
				i = i + flagsPow*icesPow-1;
			} else if (!m.isIce(p.xPosition, p.yPosition) && p.direction!=MDPState.NO_DIRECTION){
				//discard impossible states to save memory
			} else {
				//get next states, get their value, set the value in tables with the best action using 0 values
				float northVal = getVal(p, MDPState.NORTH, p.nextState(MDPState.NORTH, m, iceProbs));
				float southVal = getVal(p, MDPState.SOUTH, p.nextState(MDPState.SOUTH, m, iceProbs));
				float eastVal = getVal(p, MDPState.EAST, p.nextState(MDPState.EAST, m, iceProbs));
				float westVal = getVal(p, MDPState.WEST, p.nextState(MDPState.WEST, m, iceProbs));
				float restVal = getVal(p, MDPState.REST, p.nextState(MDPState.REST, m, iceProbs));
				byte best = MDPState.NORTH;
				float bestVal = northVal;
				if (southVal>bestVal){
					best = MDPState.SOUTH;
					bestVal = southVal;
				}
				if (eastVal>bestVal){
					best = MDPState.EAST;
					bestVal = eastVal;
				}
				if (westVal>bestVal){
					best = MDPState.WEST;
					bestVal = westVal;
				}
				if (restVal>bestVal){
					best = MDPState.REST;
					bestVal = restVal;
				}
				startTable.setValue(i, bestVal, best);
			}
		}
		for (int i = 1; i<tables.length; i++){
			System.gc();
			System.out.println("Step: " + i);
			MDPStateTable lastTable;
			MDPStateTable currentTable;
			if (remember){
				tables[i] = new MDPStateTable(stateSpace);
				lastTable = tables[i-1];
				currentTable = tables[i];
			} else {
				tables[0] = tables[1];
				lastTable = tables[0];
				tables[1] = new MDPStateTable(stateSpace);
				currentTable = tables[1];
				
			}
			for (int j=0; j<stateSpace; j++){
				MDPState p = new MDPState(j, numFlags, numIces, height, width);
				if (m.brickStatus(p.xPosition, p.yPosition)==WorldMap.WALL){	//impossible state
					j = j + flagsPow*icesPow-1; // skip all impossible states
				} else if (!m.isIce(p.xPosition, p.yPosition) && p.direction!=MDPState.NO_DIRECTION){
					//discard impossible states to save memory
				} else {
					//get next states, get their value, set the value in tables with the best action using the last table
					float northVal = getVal(p, MDPState.NORTH, p.nextState(MDPState.NORTH, m, iceProbs), lastTable);
					float southVal = getVal(p, MDPState.SOUTH, p.nextState(MDPState.SOUTH, m, iceProbs), lastTable);
					float eastVal = getVal(p, MDPState.EAST, p.nextState(MDPState.EAST, m, iceProbs), lastTable);
					float westVal = getVal(p, MDPState.WEST, p.nextState(MDPState.WEST, m, iceProbs), lastTable);
					float restVal = getVal(p, MDPState.REST, p.nextState(MDPState.REST, m, iceProbs), lastTable);
					byte best = MDPState.NORTH;
					float bestVal = northVal;
					if (southVal>bestVal){
						best = MDPState.SOUTH;
						bestVal = southVal;
					}
					if (eastVal>bestVal){
						best = MDPState.EAST;
						bestVal = eastVal;
					}
					if (westVal>bestVal){
						best = MDPState.WEST;
						bestVal = westVal;
					}
					if (restVal>bestVal){
						best = MDPState.REST;
						bestVal = restVal;
					}
					currentTable.setValue(j, bestVal, best);
				}
			}
		}
		System.out.println("Done");
		return tables;
	}
	
	public static float getVal(MDPState cur, byte action, Map<MDPState, Float> next){
		float val = 0;
		for (Entry<MDPState, Float> en: next.entrySet()){
			float weight = en.getValue();
			val = val + weight*MDPState.reward(cur, action, en.getKey());
		}
		return val;
	}
	public static float getVal(MDPState cur, byte action, Map<MDPState, Float> next, MDPStateTable t){
		float val = 0;
		for (Entry<MDPState, Float> en: next.entrySet()){
			MDPState s = en.getKey();
			float weight = en.getValue().floatValue();
			val = val + weight*(MDPState.reward(cur, action, s) + gamma*t.getValue(s.getId()));
		}
		return val;
	}
	
	private Map<Node, Integer> getIceEvidence(int i, int numIces, WorldMap map, ClimateNetwork n) {
		int[] iceSample = new int[numIces];
		int tmp = i;
		for(int j=0; j<numIces; j++){
			iceSample[j] = tmp%3;
			tmp = tmp/3;
		}
		return getEvidence(map, n, iceSample);
	}
	
	public Map<Node, Integer> getEvidence(WorldMap map, ClimateNetwork n, int[] icyPatches){
		Map<Node,Integer> ans = new HashMap<Node, Integer>();
		for (int i=0; i<icyPatches.length; i++){
			if (icyPatches[i]==MDPState.NO_ICE){
				ans.put(n.getBlockArr(map.getIceCellId(i))[0], ClimateNetwork.NO_ICE);
			} else if (icyPatches[i]==MDPState.ICE){
				ans.put(n.getBlockArr(map.getIceCellId(i))[0], ClimateNetwork.ICE);
			}
		}
		return ans;
	}
}
